package com.ruoyi.business.aspect.Redis;

import com.ruoyi.business.annotations.ratelimit.RedisRateLimit;
import com.ruoyi.common.redis.service.RedisService;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * @Description: Redis+Lua实现自定义限流
 * @Author: Wu Yuwei
 * @Date: 2023/6/28
 */
@Aspect
@Configuration
@Slf4j
public class RedisRateLimitAspect {

    @Autowired
    private RedisService redisTemplate;

    /*
        local key = "rate.limit:" .. KEYS[1]
        local limit = tonumber(ARGV[1])
        local current = tonumber(redis.call('get', key) or "0")
        if current + 1 > limit then
          return 0
        else
           -- 没有超阈值，将当前访问数量+1，并设置2秒过期（可根据自己的业务情况调整）
           redis.call("INCRBY", key,"1")
           redis.call("expire", key,"2")
           return current + 1
        end
     */
    private final String limit_lua = "local key = \"rate.limit:\" .. KEYS[1]\n" +
            "local limit = tonumber(ARGV[1])\n" +
            "local current = tonumber(redis.call('get', key) or \"0\")\n" +
            "if current + 1 > limit then\n" +
            "  return 0\n" +
            "else\n" +
            "   -- 没有超阈值，将当前访问数量+1，并设置2秒过期（可根据自己的业务情况调整）\n" +
            "   redis.call(\"INCRBY\", key,\"1\")\n" +
            "   redis.call(\"expire\", key,\"2\")\n" +
            "   return current + 1\n" +
            "end";

    @Pointcut(value = "@annotation(com.ruoyi.business.annotations.ratelimit.RedisRateLimit)")
    public void rateLimit() {

    }

    @Around("rateLimit()")
    public Object interceptor(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        RedisRateLimit rateLimit = method.getAnnotation(RedisRateLimit.class);
        if (rateLimit != null) {
            HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
            String ipAddress = getIpAddr(request);
            StringBuffer stringBuffer = new StringBuffer();
            stringBuffer.append(ipAddress).append("-")
                    .append(targetClass.getName()).append("- ")
                    .append(method.getName()).append("-")
                    .append(rateLimit.key());
            List<String> keys = Collections.singletonList(stringBuffer.toString());
            //调用lua脚本，获取返回结果，这里即为请求的次数
            DefaultRedisScript<Long> redisluaScript = new DefaultRedisScript<>(limit_lua, Long.class);
            Object number = redisTemplate.excute(
                    redisluaScript,
                    keys,
                    rateLimit.count(),
                    rateLimit.period()
            );
            if (number != null && (Long) number != 0 && (Long) number <= rateLimit.count()) {
                log.info("限流时间段内访问了第：{} 次", number.toString());
                return joinPoint.proceed();
            }
        } else {
            return joinPoint.proceed();
        }
        throw new RuntimeException("访问频率过快，被限流了");
    }

    /**
     * 获取请求的IP方法
     * @param request
     * @return
     */
    private static String getIpAddr(HttpServletRequest request) {
        String ipAddress = null;
        try {
            ipAddress = request.getHeader("x-forwarded-for");
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)) {
                ipAddress = request.getRemoteAddr();
            }
            // 对于通过多个代理的情况，第一个IP为客户端真实IP,多个IP按照','分割
            if (ipAddress != null && ipAddress.length() > 15) {
                if (ipAddress.indexOf(",") > 0) {
                    ipAddress = ipAddress.substring(0, ipAddress.indexOf(","));
                }
            }
        } catch (Exception e) {
            ipAddress = "";
        }
        return ipAddress;
    }
}
